perf(cutedsl): add Blackwell fast path to RMSNorm - #1388

Summary

This PR adds a Blackwell fast path to the CuTeDSL RMSNorm op (ops/cutedsl/ops/rms_norm.py), the only generic CuTeDSL op with no arch-specific path (e.g. swiglu.py has packed-f32x2 math and a TVM-FFI direct call; rope.py has a TMA kernel tuned on B200; and cross_entropy.py dispatches on infer_device_arch() == "blackwell", while rms_norm.py had no arch query at all).

This PR changes how the two vector fast paths — _launch_fwd_vector and _launch_bwd_fused — are launched, and the instruction mix of the backward. Their tiling, their CTA-level reduction, and the op's public API are unchanged.

1. TVM-FFI direct-call launch (host-side; the bulk of the win)

_launch_fwd_vector and _launch_bwd_fused gain a fast branch that skips the from_dlpack / memref marshalling entirely and hands PyTorch tensors straight to the compiled function:

• New _make_fwd_vector_ffi / _make_bwd_fused_ffi closures capture every constexpr, so the compiled signature is just tensors + fp32 scalars.
• New _get_fwd_vector_ffi / _get_bwd_fused_ffi compile those once per (dtype, geometry, device) against abstract tensors via make_fake_tensor, cached in a new _ffi_compile_cache. eps / offset stay runtime scalars, so changing either does not trigger a recompile.
• The stream parameter is compiled against make_fake_stream(use_tvm_ffi_env_stream=True), so it drops out of the call signature and resolves from the caller's env stream — no per-call stream query from Python.

Result: forward host cost 54us -> 15us, backward 74us -> 26us.

2. packed-f32x2 math in the fused backward (device-side)

_rms_norm_bwd_fused_vector_kernel gains a PACKED_MATH: cutlass.Constexpr branch that computes the dX/dW inner products pairwise with fma_packed_f32x2 / mul_packed_f32x2 /add_packed_f32x2 (two fp32 lanes per instruction), keeping two accumulators that fold before the CTA reduction. Selected by the new _use_packed_math, which requires sm_10x and n_cols > 4096 — the 16-warp regime where the kernel is issue-bound rather than memory-bound. Per-width measurements justifying that threshold are in its docstring.

Behavior when the fast path does not apply

• Without apache-tvm-ffi the original marshalling launch runs;
• Off data-center Blackwell (including Hopper and RTX 50xx) the scalar math runs;
• Irregular or unaligned shapes still take the existing scalar/split paths. Both halves are A/B switchable with LIGER_RMS_FORCE_NO_FFI / LIGER_RMS_FORCE_NO_PACKED.

Why the op was slow in the first place

The kernels were already faster than Triton on device time (forward 5.8us vs 7.8us at bf16 4096x2048; fused backward 26.3us vs 32.4us). Wall clock disagreed because each call spent ~48us marshalling tensors into cute.Tensor handles before it could enqueue, so the op was host-bound at every shape up to 16384x4096 — measured wall time tracked host time (~54us), not kernel time.

Results

1x B200, bf16, llama mode, op-level, cutedsl vs Triton (>1.0 = cutedsl faster):

fp32 moves the same way: forward 0.66-1.08x -> 0.99-2.29x, backward 0.80-1.73x -> 1.15-2.58x. No shape regressed in either dtype. "before" is measured by running the same build with both toggles off, so it is exactly the old code path.

The packed-math half accounts for the wide-row backward gains specifically: at bf16 16384x8192 it is 245.8us -> 222.0us (-9.7%) and at fp32 16384x6144 251.3us -> 194.6us (-22.6%), device time.

Validation

• Three new tests in test/cutedsl/test_rms_norm.py cover ground the existing parity suite cannot reach — it only goes up to 4096 wide (so it never compiles the packed backward) and always runs on the default stream:

  • test_rms_norm_ffi_matches_marshalling_path — bit-identity between the two launch paths, affine and non-affine; the non-affine case is what pins the dummy-W layout.
  • test_rms_norm_packed_backward_matches_scalar — packed vs scalar backward at 8192.
  • test_rms_norm_runs_on_the_current_stream — stream binding via CUDA graph capture.

  Both comparison tests assert the two variants were really compiled, so they cannot degenerate into comparing a path against itself. Mutation-checked: corrupting the packed dW accumulation or the non-affine dummy-W layout makes them fail.

• Full cutedsl suite on B200: 579 passed with an identical 33-failure set before and after (pre-existing bf16/fp32 parity gaps in cross_entropy and RMSNorm's bf16-llama dW) — no new failures.

• ruff check / ruff format clean.

Packaging note

apache-tvm-ffi is added to the cutedsl extra. It stays optional at runtime, but the entire forward win depends on it and the existing swiglu kernel already has the same latent dependency. Happy to drop this if you would rather keep the extra minimal.